import timm
import torch
from torch import nn


class ISICModel(nn.Module):
    def __init__(self, model_name, model_path=None, drop_path_rate=0, drop_rate=0, pretrained=True, checkpoint_path=None):
        super(ISICModel, self).__init__()
        self.model = timm.create_model(
            model_name, 
            pretrained=pretrained, 
            pretrained_cfg_overlay=dict(file=model_path) if model_path is not None else None,
            checkpoint_path=checkpoint_path,
            num_classes=0,
            drop_rate=drop_rate, 
            drop_path_rate=drop_path_rate)
        
        in_features = self.model.head_hidden_size
        self.head_2018 = nn.Linear(in_features, 7)
        self.head_2019 = nn.Linear(in_features, 9)
        self.head_2020 = nn.Linear(in_features, 1)
        self.head_2024 = nn.Linear(in_features, 1)

    def forward(self, images):
        feature = self.model(images)
        return {
            2018: nn.Softmax()(self.head_2018(feature)),
            2019: nn.Softmax()(self.head_2019(feature)),
            2020: nn.Sigmoid()(self.head_2020(feature)).squeeze(),
            2024: nn.Sigmoid()(self.head_2024(feature)).squeeze(),
        }


class ISICModelEdgnet(nn.Module):
    def __init__(self, model_name, model_path=None, drop_path_rate=0, drop_rate=0, pretrained=True, checkpoint_path=None):
        super(ISICModelEdgnet, self).__init__()
        self.model = timm.create_model(
            model_name, 
            global_pool='avg',
            pretrained=pretrained, 
            pretrained_cfg_overlay=dict(file=model_path) if model_path is not None else None,
            checkpoint_path=checkpoint_path,
            num_classes=0,)
        
        in_features = self.model.head_hidden_size
        self.head_2018 = nn.Linear(in_features, 7)
        self.head_2019 = nn.Linear(in_features, 9)
        self.head_2020 = nn.Linear(in_features, 1)
        self.head_2024 = nn.Linear(in_features, 1)
        
    def forward(self, images):
        feature = self.model(images)
        return {
            2018: nn.Softmax()(self.head_2018(feature)),
            2019: nn.Softmax()(self.head_2019(feature)),
            2020: nn.Sigmoid()(self.head_2020(feature)).squeeze(),
            2024: nn.Sigmoid()(self.head_2024(feature)).squeeze(),
        }



def setup_model(model_name, model_path=None, checkpoint_path=None, num_classes=1, drop_path_rate=0, drop_rate=0, 
                device: str = 'cuda', model_maker=ISICModel):
    model = model_maker(model_name, model_path=model_path, pretrained=True, checkpoint_path=checkpoint_path, 
                       drop_path_rate=drop_path_rate, drop_rate=drop_rate)

    return model.to(device)

